# -*- coding:utf-8 -*-

# the cgraph class maintains a graph's topology and message passing schedule.
# IMPORTANT ensure every varnode and factnode has unique name and the name of varnode and factnode should be diiferent

from  builtins import range
from functools import reduce
import numpy as np
from  cnode import *
from cmsg import *
from motionModel import *
from mpl_toolkits.mplot3d import Axes3D
#from hybridNode import *
import logging



class cGraph(object):
    def __init__(self,name = None):
        if name is None:
            name = 'agraph'
        self.name = name
        self.connections = list([]) # key pairs represents the fact-var connection

        self.nodes    = list([])
        self.nodeNames = dict({}) # key is the name and the value is node Inde 
        
        self.agentNodes =list([])  # list of the agent node in the nodes, represented by index


    def getNodeIndex(self,node):
        '''
            get index givin node (string or instance of cNode)
        '''
        if(isinstance(node,basestring)):
            if(self.nodeNames.has_key(node)):
                return self.nodeNames[node]
            else:
                return None
        elif isinstance(node,int):
            if node>=0 and node< len(self.nodeNames):
                return node
            else:
                raise Exception('out of index')
        elif(isinstance(node, cNode)):
            raise Exception('Don\'t support this kind of data')
        else:
            raise Exception('Don\'t support this kind of data')

    
    def getNodeInstance(self,node):
        '''
            get the varnode instance given varIndex
        '''
        if isinstance(node,list):
            nodesR  =  [] 
            for node1 in node:
                if isinstance(node1,str):
                    node1 = self.getNodeIndex(node1)
                nodesR.append( self.nodes[node1] )   
            return nodesR
        else:
            if isinstance(node,str):
                node = self.getNodeIndex(node)
            return self.nodes[node]

    def addNode(self,node):
        nodeName   = node.name
        assert(self.nodeNames.has_key(nodeName)==False)
        self.nodes.append(node)
        nodeIndex  = len(self.nodes)-1
        self.nodeNames[nodeName] = nodeIndex
        # self.agentNodes
        if isinstance(node,agentVarNode):
           self.agentNodes.append(nodeIndex)
        return nodeIndex
    # add var nodes to the graph
    def addNodes(self,nodes):
        if(isinstance(nodes,list)==False):
            nodes = [nodes]
        for node in nodes:
            self.addNode(node)

    def connectVarNodes(self,varnodes,factnode,connectionTypes='bi'):
        '''
            connect existing varnodes with an existing factnode, and will insert into the graph
            varnode can be a list, basestring, integer, object of a varnode
            factnode MUST be an instance of a factNode
            connectionTypes: connection types of varnodes, can be a string or a list
        '''
        factNodeAlreadyIn  =  False
        if(isinstance(factnode,str)):
            factIndex= self.getNodeIndex(factnode)
            factnode = self.getNodeInstance(factnode)
            factNodeAlreadyIn = True

        if(isinstance(factnode,factNode)==False):
            raise Exception('factnode should be an instance of factnode')
        if factNodeAlreadyIn==False:
            factIndex = self.addNode(factnode)  # TODO ERROR This should be delete
       
        if(isinstance(varnodes,list)==False):
            varnodes = [varnodes]
        if(isinstance(connectionTypes,list)==False):
            connectionTypes = [connectionTypes]

        assert(len(connectionTypes)==len(varnodes))

        for index in range(0,len(varnodes)):
            varnode = varnodes[index]
            varIndex = varnode
            if(isinstance(varIndex,int)==False):
                varIndex = self.getNodeIndex(varnode)
            varnode  = self.getNodeInstance(varIndex)
            self.connections.append({factIndex,varIndex})
            # some operation in the node
            varnode.addAdjNodes(factnode,connectionTypes[index])

    # disconnect a varnode from factnode
    # the corresponding connection between the factnode and other varnode will also be disconnected 
    def disconnectNodes(self,varnode,factnode):
        pass

    def delFactNode(self,factnodes):
        pass

    def delVarNode(self,varnodes):
        pass

    
    # message passing related functions 
    
    def passMsg(self,fromnode,tonode=None):
        '''
            pass Msg from fromnode to tonode, following those instruction:
            Ensure two nodes are node string/instance
            
            If node2 is none, then node1 will broadcast to all the out/bi ports
            TODO check the connectivity of node1 and node2 
        '''
        if tonode!=None:
            fromnode = self.getNodeInstance(fromnode)
            tonode = self.getNodeInstance(tonode)
            logging.debug( "++++++++\n Sending data:" + fromnode.name + '->'+ tonode.name  + "-------> ")
            port = tonode.getAdjNodeIndex(fromnode)
            portTo = fromnode.getAdjNodeIndex(tonode)
            tonode.receiveMessage(fromnode.sendMsg(portTo),port)
            logging.debug( '<<<<<<<------------\n\n')
    def flushMsgs(self):
        '''
            flush messages of all nodes
        '''
        for node in self.nodes:
            node.flushMsgs()
    def flushConnections(self):
        '''
            flush all of the connections
        '''
        for node in self.nodes:
            node.flushConnections()

if __name__ == '__main__':
    # logging.basicConfig(level=logging.DEBUG,format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s')

    logging.basicConfig(level=logging.ERROR,format='%(message)s')
    logging.debug('Hello')

    np.random.seed(123)

    # simulation parameters
    ps ={}
    # agent nodes amount 
    ps['agentNodes'] = 10

    # dimension settting
    ps['dim']        = 3

    # transform noise level, sigma: standard derivation
    ps['xTransformNoise'] = np.eye(ps['dim'])   #np.diag([1,1])
    ps['satMeasureNoise'] = np.array([1.0],dtype=float)
    ps['mutualMeasureNoise']  = np.array([1.0],dtype=float)
    # satllite amount avaliable
    ps['satNum']  = 10
    
    # connection Rate, from 0 to 1, what is connection rate?
    ps['conRate'] = 0.5

    xPosHistory = np.zeros((ps['dim'],ps['satNum'],2))

    # satllite connection amount for each agent 
    satAmount = [3,4,5,4,2,3,5,2,3,4,2,1,2]


    # randomly determine the satellite connection with nodes
    # satConnectionIndex, list of list of the satellite connection for each agent
    satConnectionIndex =[]
    for satNum in satAmount:
        satIndex     = []
        tempsatIndex = np.floor(np.random.uniform(0,ps['satNum'],satNum))
        for index in tempsatIndex:
            satIndex.append(int(index))
        satConnectionIndex.append(satIndex)

    # randomly generate satPoses (satllite positions)
    satPoses  = np.random.uniform(-1000,1000,(ps['dim'],ps['satNum']))
    
    # randomly generate nodes' true positions
    xTrueLast = np.random.uniform(-100,100,(ps['dim'],ps['agentNodes']))
    xPosHistory[:,:,0] = xTrueLast
    
    # randomly generate cooperative connections among nodes
    # CoopConnections: a list of cooperative ALL POSSIBLE pairs, same structure with satConnectionIndex
    CoopConnections= []
    for i in range(ps['agentNodes']):
        for j in range(0,i):
            CoopConnections.append([i,j])
    
    # and pick some of them
    connectionAmount = int(np.ceil(ps['conRate']*len(CoopConnections)))
    connectionIndexes = np.random.choice(range(len(CoopConnections)),connectionAmount,False)

    '''
        Finnally, we can generate the factor graph
    '''

    # create a graph and its topology
    g =  cGraph('correction')

    varnodes =  []
    for i  in range(ps['agentNodes']):
        varnode = priorFactNode('priorX'+str(i))
        priorMsg = normalDistributeMsg('priorX'+str(i)+'msg',xTrueLast[:,i], np.eye(ps['dim']))
        varnode.setPriorMsg(priorMsg)
        varnodes.append(varnode)
        varnode = agentVarNode('X'+str(i))
        varnodes.append(varnode)

    g.addNodes(varnodes)

    #true transformation and measurement generating 
    xTrueLast = xTrueLast

    # satMeasurements for each connection
    satMeasurements = list()
    for j in range(ps['agentNodes']):
        agentPos =  xTrueLast[:,j]
        satMeasureforSingleAgent =[] 
        for satIndex  in satConnectionIndex[j]:
            satPos = satPoses[:,satIndex]
            dist   = np.linalg.norm(agentPos-satPos)
            measure = dist+np.random.normal(0,ps['satMeasureNoise'])
            satMeasureforSingleAgent.append(measure)
        satMeasurements.append(satMeasureforSingleAgent)

    # set measurement of Coop connections
    ranges  =  [] 
    for connectionIndex in connectionIndexes:
        node1Cordinate  = xTrueLast[:,CoopConnections[connectionIndex][0]]
        node2Cordinate  = xTrueLast[:,CoopConnections[connectionIndex][1]]
        ranges.append(np.linalg.norm(node1Cordinate -node2Cordinate)+np.random.normal(0,ps['mutualMeasureNoise']))



    for i  in range(ps['agentNodes']):

        factnode = mobilityFactNode('x'+str(i)+'motion')

        motionModel = randomWalkMotion({'noiseE':np.zeros(ps['dim']),'noiseCov':ps['xTransformNoise']})

        factnode.setMobilityModel(motionModel)
        
        g.connectVarNodes(['priorX'+str(i),'X'+str(i)],factnode,['out','in'])
        
        for j  in range(satAmount[i]):
            satNode = satelliteFactNode('sat'+str(i)+'_'+str(j))
            satMsg  = vMsg('sat'+str(i)+'_'+str(j)+'msg',
                      satPoses[:,satConnectionIndex[i][j]],ps['satMeasureNoise'],
                      satMeasurements[i][j])
            satNode.setPriorMsg(satMsg)

            g.connectVarNodes('X'+str(i),satNode,'in')
        
    index = 0
    for connectionIndex in connectionIndexes:
         node1Index  = CoopConnections[connectionIndex][0]
         node2Index  = CoopConnections[connectionIndex][1]
         mutualConnectionFact  =  mutualConNode('mutual'+str(node1Index)+'_'+str(node2Index),ranges[index],ps['mutualMeasureNoise']**2)
         index+=1
         # TODO assign measurement to this mutualConnectionFact
         g.connectVarNodes(['X'+str(node1Index),'X'+str(node2Index)],mutualConnectionFact,['bi','bi'])

        

    # The following code realizes HSPAWN algorithm

    #first, we get the priori distribution
    for agentIndex in range(ps['agentNodes']):
        g.passMsg('priorX'+str(agentIndex),'x'+str(agentIndex)+'motion')
        g.passMsg('x'+str(agentIndex)+'motion','X'+str(agentIndex))
    
    #Secoondly, we send all satellite positions
    for agentIndex in range(ps['agentNodes']):
        for j in range(satAmount[agentIndex]):
            satName =  'sat'+str(agentIndex)+'_'+str(j)
            g.passMsg(satName,'X'+str(agentIndex))

    # then start p2p message passing
    iterTime =2
    for iterate in range(iterTime):
        for connectionIndex in connectionIndexes:
            node1Index  = CoopConnections[connectionIndex][0]
            node2Index  = CoopConnections[connectionIndex][1]
            # TODO it must transmit all the information to factnode and process. we should change this
            g.passMsg('X'+str(node1Index),'mutual'+str(node1Index)+'_'+str(node2Index))
            g.passMsg('X'+str(node2Index),'mutual'+str(node1Index)+'_'+str(node2Index))

            g.passMsg('mutual'+str(node1Index)+'_'+str(node2Index),'X'+str(node1Index))
            g.passMsg('mutual'+str(node1Index)+'_'+str(node2Index),'X'+str(node2Index))      
            
            #TODO refresh the inMsgs

    # plot results
    agentNodes = g.getNodeInstance(g.agentNodes)

    satPosesTmp = np.zeros((ps['dim'],ps['satNum']))
    index= 0
    for agentNode in agentNodes:
        satPosesTmp[:,index]  = agentNode.beliefs.mean
        index +=1
    
    xPosHistory[:,:,1] = satPosesTmp

    fig= plt.figure(1)
    ax = fig.gca(projection='3d')

    for i in range(xPosHistory.shape[1]):
        h = ax.plot(xPosHistory[0,i,:],xPosHistory[1,i,:],xPosHistory[2,i,:])
        plt.setp(h,marker='+',ms=30,mew=3)


    # h = plt.plot(xPosHistory[0,:,0],xPosHistory[1,:,0],'+',ms=30,mew=3)


    plt.show()